from CLaplace5 import *
from matplotlib.pyplot import plot, savefig, title, legend, ylim, cla, xlabel, ylabel, annotate
#from sklearn.utils.extmath import fast_dot
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm
from time import clock

def Plot_Result(Z, Nx, Ny, Xmax, Xmin, Ymax, Ymin):
	Border_Up, Border_Down, Border_Left, Border_Right = Border_Generator(Nx, Ny, Border_Type="matrix")
	xs = np.linspace(Xmin, Xmax, Nx+2)
	ys = np.linspace(Ymin, Ymax, Ny+2)
	X, Y = np.meshgrid(xs, ys)
	Z = np.reshape(Z, [Ny,Nx],"F")#reshape the matrix as Fortran
	Z = np.bmat([[Border_Left, np.bmat([[Border_Up], [Z], [Border_Down]]), Border_Right]])
	Z = np.array(Z)
	fig = plt.figure()
	ax = fig.gca(projection='3d')

	cset = ax.contourf(X, Y, Z, zdir='z', offset=0, cmap=cm.coolwarm)
	cset = ax.contourf(X, Y, Z, zdir='x', offset=-1, cmap=cm.coolwarm)
	cset = ax.contour(X, Y, Z, zdir='y', offset=60, cmap=cm.coolwarm)
	surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, alpha=0.5, cmap=cm.coolwarm, linewidth=0.01)

	fig.colorbar(surf, shrink=0.5, aspect=5)
	ax.set_xlabel('X')
	ax.set_xlim(Xmin, Xmax*1.05)
	ax.set_ylabel('Y')
	ax.set_ylim(Ymin, Ymax*1.05)
	ax.set_zlabel('Z')
	ax.set_zlim(np.min(Z), np.max(Z)*1.05)
	plt.show()

for N in range(4,62,2):
	print(N,end="",flush=True)
	A_shape, b, D, H = Generate_Matrix_A_b(N, N, 5, 0, 5, 0)
	print("\t--->\t",end="",flush=True)
	start_time = clock()
	Z, Error, count = Method_Iterative(A_shape, b, D, H, 1e-5)
	end_time = clock()
	with open("using_cython_without_modifing_python_code", "a") as benchmark_file:
		print(N, end_time-start_time, file=benchmark_file)
	print("Done!",flush=True)
